# _*_ coding: utf-8 _*_
import pickle
from random import sample, uniform

from tqdm import trange

from Encodings import MultiToursSolution


def copy(target):
    return pickle.loads(pickle.dumps(target))


class ABC:  # Artificial Bee Colony
    def __init__(self, data):
        self.name = __class__.__name__
        # -- read instance from instance data --
        self.data = copy(data)
        self.V = self.data.cities  # city list
        self.n = self.data.n  # number of cities
        self.m = self.data.m  # number of salesmen
        self.A = self.data.A  # accessibility matrix
        self.Depot = self.data.depot  # depot
        self.V.remove(self.Depot)  # exclude depot
        self.Distance = self.data.distance_matrix  # distance matrix
        self.tour_cost = self.data.cost  # cost function for a single tour
        self.access = self.data.access  # fast computation access func
        # -- parameters ---
        self.NP = 10  # the number of employed bees
        self.NO = self.NP  # the number of onlooker bees
        self.Pc = 0.8  # deepcopy probability
        self.Pc_min = 0  # min deepcopy probability
        self.Pc_max = 0.8  # max deepcopy probability
        self.Pbts = 0.8  # the binary tournament selection probability
        self.limit_scout = self.NP
        # -- solution initialization --
        self.Population = [MultiToursSolution().init_solution(cities=self.V, distance=self.Distance, access=self.access,
                                                              depot=self.Depot, m=self.m)
                           for _ in range(self.NP)]
        self.optima = copy(min(self.Population, key=lambda x: x.total))
        self.evo_list = []
        # -- auxiliary --
        self.entropy = []  # assignment entropy to measure the diversity of population
        self.Pc_list = []  # Pc value in history
    
    def optimize(self, iteration=1000, trial=1):
        pbar = trange(iteration)
        pbar.set_description(f'{self.name} on {self.data.name} trial {trial}')
        for it in pbar:
            # update Pc value
            self.Pc = (self.Pc_max - self.Pc_min) * it / iteration + self.Pc_min
            self.Pc_list.append(self.Pc)  # record Pc value
            # employed bees phase
            for i in range(self.NP):
                # generate a neighbor solution
                solution = self.Population[i].neighborhood(copy_prob=self.Pc, distance=self.Distance,
                                                           cost_func=self.tour_cost, access=self.access)
                if solution.total < self.Population[i].total:
                    self.Population[i] = copy(solution)
                    self.Population[i].stags = 0
                else:
                    self.Population[i].stags += 1
            # onlooker bees phase
            for i in range(self.NO):
                # use binary tournament to select some population
                tournament = sorted(sample(list(range(self.NP)), k=2), key=lambda x: self.Population[x].total)
                p = tournament[0] if uniform(0, 1) < self.Pbts else tournament[1]
                # generate a neighbor solution
                solution = self.Population[p].neighborhood(copy_prob=self.Pc, distance=self.Distance,
                                                           cost_func=self.tour_cost, access=self.access)
                if solution.total < self.Population[i].total:
                    self.Population[i] = copy(solution)
                    self.Population[i].stags = 0
                else:
                    self.Population[i].stags += 1
            # update optima & record evolutionary history
            if self.optima.total > min(self.Population, key=lambda x: x.total).total:
                self.optima = copy(min(self.Population, key=lambda x: x.total))
            self.evo_list.append(self.optima.total)
            pbar.set_postfix_str(f'best: {self.optima.total:.2f}')
            # scout bees phase
            for i in range(self.NP):
                if self.Population[i].stags > self.limit_scout:
                    self.Population[i] = MultiToursSolution().init_solution(cities=self.V, distance=self.Distance,
                                                                            access=self.access,
                                                                            depot=self.Depot, m=self.m)
        return self
